#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HAI服务器管理脚本
用于快速管理HAI优化系统的服务器实例
"""

import requests
import json
import time
import sys
import argparse
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Any

class HAIServerManager:
    def __init__(self, base_url: str = "http://localhost:8080"):
        self.base_url = base_url
        
    def get_system_status(self) -> Dict[str, Any]:
        """获取系统状态"""
        try:
            response = requests.get(f"{self.base_url}/api/v1/system/status", timeout=10)
            if response.status_code == 200:
                return response.json()
            else:
                print(f"❌ 获取系统状态失败: {response.status_code}")
                return {}
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return {}
    
    def get_servers(self) -> List[Dict[str, Any]]:
        """获取所有服务器列表"""
        try:
            response = requests.get(f"{self.base_url}/api/v1/servers", timeout=10)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    return data.get('data', [])
            else:
                print(f"❌ 获取服务器列表失败: {response.status_code}")
                return []
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return []
    
    def destroy_server(self, instance_id: str) -> bool:
        """销毁指定服务器"""
        try:
            response = requests.delete(f"{self.base_url}/api/v1/servers/{instance_id}", timeout=60)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    print(f"✅ 服务器 {instance_id} 销毁成功")
                    return True
                else:
                    print(f"❌ 服务器 {instance_id} 销毁失败: {data.get('error')}")
                    return False
            else:
                print(f"❌ 销毁请求失败: {response.status_code}")
                return False
        except Exception as e:
            print(f"❌ 销毁异常: {e}")
            return False
    
    def create_server(self, region: str = "ap-beijing", gpu_type: str = "basic") -> str:
        """创建新服务器"""
        try:
            response = requests.post(f"{self.base_url}/api/v1/servers", 
                                   json={'region': region, 'gpu_type': gpu_type}, 
                                   timeout=60)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    instance_id = data.get('instance_id')
                    print(f"✅ 服务器创建成功: {instance_id}")
                    return instance_id
                else:
                    print(f"❌ 服务器创建失败: {data.get('error')}")
                    return ""
            else:
                print(f"❌ 创建请求失败: {response.status_code}")
                return ""
        except Exception as e:
            print(f"❌ 创建异常: {e}")
            return ""
    
    def create_servers_parallel(self, count: int) -> int:
        """并行创建多个服务器"""
        created_count = 0
        created_servers = []
        
        def create_single_server(index: int) -> str:
            """创建单个服务器"""
            try:
                print(f"🚀 创建服务器 {index+1}/{count}")
                response = requests.post(f"{self.base_url}/api/v1/servers", 
                                       json={'region': 'ap-beijing', 'gpu_type': 'basic'}, 
                                       timeout=60)
                if response.status_code == 200:
                    data = response.json()
                    if data.get('success'):
                        instance_id = data.get('instance_id')
                        print(f"✅ 服务器创建成功: {instance_id}")
                        return instance_id
                    else:
                        print(f"❌ 服务器创建失败: {data.get('error')}")
                        return ""
                else:
                    print(f"❌ 创建请求失败: {response.status_code}")
                    return ""
            except Exception as e:
                print(f"❌ 创建异常: {e}")
                return ""
        
        # 使用线程池并行创建 - 增加并发数支持大规模部署
        with ThreadPoolExecutor(max_workers=min(count, 20)) as executor:
            # 提交所有创建任务
            future_to_index = {
                executor.submit(create_single_server, i): i 
                for i in range(count)
            }
            
            # 收集结果
            for future in as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    instance_id = future.result()
                    if instance_id:
                        created_count += 1
                        created_servers.append(instance_id)
                except Exception as e:
                    print(f"❌ 创建任务 {index+1} 异常: {e}")
        
        return created_count
    
    def destroy_servers_parallel(self, servers_to_destroy: List[Dict]) -> int:
        """并行销毁多个服务器"""
        destroyed_count = 0
        
        def destroy_single_server(server_info: Dict, index: int) -> bool:
            """销毁单个服务器"""
            try:
                instance_id = server_info.get('instance_id')
                print(f"🗑️ 销毁服务器 {index+1}/{len(servers_to_destroy)}: {instance_id}")
                
                response = requests.delete(f"{self.base_url}/api/v1/servers/{instance_id}", 
                                        timeout=60)
                if response.status_code == 200:
                    data = response.json()
                    if data.get('success'):
                        print(f"✅ 服务器销毁成功: {instance_id}")
                        return True
                    else:
                        print(f"❌ 服务器销毁失败: {data.get('error')}")
                        return False
                else:
                    print(f"❌ 销毁请求失败: {response.status_code}")
                    return False
            except Exception as e:
                print(f"❌ 销毁异常: {e}")
                return False
        
        # 使用线程池并行销毁 - 增加并发数支持大规模销毁
        with ThreadPoolExecutor(max_workers=min(len(servers_to_destroy), 20)) as executor:
            # 提交所有销毁任务
            future_to_index = {
                executor.submit(destroy_single_server, server, i): i 
                for i, server in enumerate(servers_to_destroy)
            }
            
            # 收集结果
            for future in as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    success = future.result()
                    if success:
                        destroyed_count += 1
                except Exception as e:
                    print(f"❌ 销毁任务 {index+1} 异常: {e}")
        
        return destroyed_count
    
    def scale_to_target(self, target_count: int) -> bool:
        """缩放到目标服务器数量"""
        print(f"🎯 开始缩放到 {target_count} 个服务器...")
        
        # 获取当前服务器列表
        servers = self.get_servers()
        if servers is None:
            print("❌ 无法获取服务器列表")
            return False
        
        current_count = len(servers)
        print(f"📊 当前服务器数量: {current_count}")
        
        if current_count == target_count:
            print(f"✅ 服务器数量已经是 {target_count}，无需调整")
            return True
        
        if current_count > target_count:
            # 需要销毁服务器
            destroy_count = current_count - target_count
            print(f"🗑️ 需要销毁 {destroy_count} 个服务器")
            
            # 并行销毁多余的服务器
            print(f"🗑️ 并行销毁 {destroy_count} 个服务器...")
            destroyed = self.destroy_servers_parallel(servers[:destroy_count])
            print(f"✅ 成功销毁 {destroyed} 个服务器")
            
        else:
            # 需要创建服务器
            create_count = target_count - current_count
            print(f"🚀 需要创建 {create_count} 个服务器")
            
            # 并行创建服务器
            print(f"🚀 并行创建 {create_count} 个服务器...")
            created = self.create_servers_parallel(create_count)
            print(f"✅ 成功创建 {created} 个服务器")
        
        # 等待系统稳定
        print("⏳ 等待系统稳定...")
        time.sleep(5)
        
        # 验证结果
        final_servers = self.get_servers()
        final_count = len(final_servers) if final_servers else 0
        print(f"📊 最终服务器数量: {final_count}")
        
        if final_count == target_count:
            print(f"✅ 缩放完成！当前有 {final_count} 个服务器")
            return True
        else:
            print(f"⚠️ 缩放未完全成功，当前有 {final_count} 个服务器，目标 {target_count} 个")
            return False
    
    def show_status(self):
        """显示系统状态"""
        print("📊 HAI优化系统状态:")
        print("=" * 50)
        
        # 系统状态
        status = self.get_system_status()
        if status and status.get('success'):
            data = status['data']
            print(f"系统运行: {data['system']['running']}")
            print(f"扩缩容运行: {data['scaling']['running']}")
            print(f"总服务器: {data['servers']['servers']['total']}")
            print(f"活跃服务器: {data['servers']['servers']['available']}")
            print(f"队列长度: {data['queue']['total_tasks_in_queue']}")
            print()
        
        # 服务器列表
        servers = self.get_servers()
        if servers:
            print(f"📋 服务器列表 ({len(servers)} 个):")
            for i, server in enumerate(servers, 1):
                instance_id = server.get('instance_id', 'Unknown')
                status = server.get('status', 'Unknown')
                region = server.get('region', 'Unknown')
                gpu_type = server.get('gpu_type', 'Unknown')
                print(f"  {i:2d}. {instance_id} | {status} | {region} | {gpu_type}")
        else:
            print("❌ 无法获取服务器列表")
    
    def destroy_all(self):
        """销毁所有服务器"""
        print("🗑️ 销毁所有服务器...")
        return self.scale_to_target(0)
    
    def keep_one(self):
        """只保留1个服务器"""
        print("🎯 只保留1个服务器...")
        return self.scale_to_target(1)

def main():
    parser = argparse.ArgumentParser(description='HAI服务器管理脚本')
    parser.add_argument('action', choices=['status', 'scale', 'destroy-all', 'keep-one'], 
                       help='操作类型')
    parser.add_argument('--target', type=int, default=1, 
                       help='目标服务器数量 (仅用于scale操作)')
    parser.add_argument('--url', default='http://localhost:8080', 
                       help='HAI优化系统URL')
    
    args = parser.parse_args()
    
    manager = HAIServerManager(args.url)
    
    if args.action == 'status':
        manager.show_status()
    elif args.action == 'scale':
        manager.scale_to_target(args.target)
    elif args.action == 'destroy-all':
        manager.destroy_all()
    elif args.action == 'keep-one':
        manager.keep_one()

if __name__ == '__main__':
    main()
